package net.cyue.ort.llm.generator;

import net.cyue.ort.llm.ModelAdapter;
import net.cyue.ort.llm.ModelException;
import net.cyue.ort.llm.cache.CacheManager;
import net.cyue.ort.llm.data.ModelInput;
import net.cyue.ort.llm.data.ModelOutput;
import net.cyue.ort.llm.data.Token;
import net.cyue.ort.llm.sampling.SamplingConfig;
import net.cyue.ort.llm.sampling.SamplingStrategy;
import net.cyue.ort.llm.sampling.SamplingUtils;
import net.cyue.ort.llm.util.TokenManager;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
import java.util.List;
import java.util.Map;
import java.util.Objects;

/**
 * Beam Search 文本生成器
 * 基于束搜索策略生成高概率序列
 */
public class BeamSearchTextGenerator implements TextGenerator {
    private static final Logger logger = LoggerFactory.getLogger(BeamSearchTextGenerator.class);

    private final ModelAdapter model;
    private final TokenManager tokenManager;
    private final CacheManager cacheManager;
    private final SamplingStrategy samplingStrategy;

    public BeamSearchTextGenerator(
        ModelAdapter model,
        TokenManager tokenManager,
        CacheManager cacheManager,
        SamplingStrategy samplingStrategy
    ) {
        this.model = Objects.requireNonNull(model, "model must not be null");
        this.tokenManager = Objects.requireNonNull(tokenManager, "tokenManager must not be null");
        this.cacheManager = Objects.requireNonNull(cacheManager, "cacheManager must not be null");
        this.samplingStrategy = Objects.requireNonNull(samplingStrategy, "samplingStrategy must not be null");
    }

    @Override
    public String generate(String prompt, GenerationConfig config, GenerationCallback callback) {
        long startTime = System.currentTimeMillis();
        if (logger.isDebugEnabled()) {
            logger.debug("╔═══════════════════════════════════════════════════════════════╗");
            logger.debug("║          BeamSearchTextGenerator - 开始生成                  ║");
            logger.debug("╚═══════════════════════════════════════════════════════════════╝");
            logger.debug("📝 Prompt (长度: {}): {}", prompt.length(),
                prompt.length() > 100 ? prompt.substring(0, 100) + "..." : prompt);
            logger.debug("⚙️  配置参数:");
            logger.debug("   • maxNewTokens: {}", config.getMaxNewTokens());
            logger.debug("   • numBeams: {}", config.getNumBeams());
            logger.debug("   • beamDepth: {}", config.getBeamDepth());
            logger.debug("   • lengthPenalty: {}", config.getLengthPenalty());
        }

        List<Long> promptTokens = tokenManager.tokenize(prompt);
        if (logger.isDebugEnabled()) {
            logger.debug("🔤 Prompt Tokens: size={}, preview={}",
                promptTokens.size(), previewIds(promptTokens, 12));
        }
        Map<String, Object> initialCache = cacheManager.createEmptyCache(1, promptTokens.size(), model.getConfig());
        SamplingConfig samplingConfig = SamplingConfig.fromGenerationConfig(config);
        SimpleCallbackState state = new SimpleCallbackState(callback);

        List<BeamState> beams = new ArrayList<>();
        beams.add(BeamState.initial(initialCache));

        String result;
        try {
            result = runBeamSearch(promptTokens, samplingConfig, config, beams, state);
            state.finish();
        } catch (ModelException e) {
            state.error(e);
            logger.error("Beam search 生成失败", e);
            throw new GenerationException("Beam search generation failed", e);
        } finally {
            for (BeamState beam : beams) {
                cacheManager.cleanupCache(beam.cache());
            }
        }

        long duration = System.currentTimeMillis() - startTime;
        if (logger.isDebugEnabled()) {
            logger.debug("");
            logger.debug("╔═══════════════════════════════════════════════════════════════╗");
            logger.debug("║          BeamSearchTextGenerator - 生成完成                  ║");
            logger.debug("╚═══════════════════════════════════════════════════════════════╝");
            logger.debug("📊 统计信息:");
            logger.debug("   • 总耗时: {}ms ({}s)", duration, String.format("%.2f", duration / 1000.0));
            logger.debug("   • 输出长度: {} 字符", result.length());
        }
        return result;
    }

    private String runBeamSearch(
        List<Long> promptTokens,
        SamplingConfig samplingConfig,
        GenerationConfig config,
        List<BeamState> beams,
        SimpleCallbackState state
    ) throws ModelException {
        String bestText = "";
        int beamWidth = Math.max(1, config.getNumBeams());
        float lengthPenalty = Math.max(0f, config.getLengthPenalty());
        List<Long> committedTokens = new ArrayList<>();

        for (int step = 0; step < config.getMaxNewTokens(); step++) {
            List<BeamState> expanded = new ArrayList<>();
            boolean allFinished = true;

            for (BeamState beam : beams) {
                if (beam.isFinished()) {
                    expanded.add(beam);
                    continue;
                }
                allFinished = false;

                List<Long> context = buildContext(promptTokens, beam);
                ModelInput input = buildModelInput(context, beam);
                ModelOutput output = model.runInference(input);
                float[] logits = output.getLastTokenLogits();
                Map<String, Object> cacheSnapshot = cacheManager.extractFromResult(output);

                List<TokenScore> candidates = selectCandidates(logits, context, samplingConfig, beamWidth);
                if (candidates.isEmpty()) {
                    cacheManager.cleanupCache(cacheSnapshot);
                    continue;
                }

                List<Map<String, Object>> caches = allocateCaches(cacheSnapshot, candidates.size());
                for (int idx = 0; idx < candidates.size(); idx++) {
                    TokenScore candidate = candidates.get(idx);
                    Map<String, Object> childCache = caches.get(idx);
                    BeamState child = beam.extend(candidate, childCache, tokenManager.isEndOfSequence(candidate.tokenId()));
                    expanded.add(child);
                }

                cacheManager.cleanupCache(beam.cache());
            }

            if (expanded.isEmpty() || allFinished) {
                break;
            }

            expanded.sort(Comparator.comparingDouble((BeamState b) -> b.normalizedScore(lengthPenalty)).reversed());
            List<BeamState> nextBeams = new ArrayList<>(Math.min(beamWidth, expanded.size()));
            for (int i = 0; i < expanded.size(); i++) {
                BeamState stateToHandle = expanded.get(i);
                if (i < beamWidth) {
                    nextBeams.add(stateToHandle);
                } else {
                    cacheManager.cleanupCache(stateToHandle.cache());
                }
            }
            beams.clear();
            beams.addAll(nextBeams);

            int stableLen = stablePrefixLength(beams);
            if (stableLen > committedTokens.size() && !beams.isEmpty()) {
                List<Long> reference = beams.get(0).generatedTokens();
                List<Long> incrementalTokens = reference.subList(committedTokens.size(), stableLen);
                String deltaText = tokenManager.detokenize(new ArrayList<>(incrementalTokens));
                if (deltaText != null && !deltaText.isEmpty()) {
                    state.appendText(deltaText);
                }
                committedTokens.addAll(incrementalTokens);
            }

            if (logger.isDebugEnabled()) {
                logStepSummary(step, beams, stableLen, committedTokens);
            }
        }

        BeamState bestBeam = beams.stream()
            .max(Comparator.comparingDouble(b -> b.normalizedScore(lengthPenalty)))
            .orElse(null);

        if (bestBeam != null && !bestBeam.generatedTokens().isEmpty()) {
            int shared = commonPrefixLength(committedTokens, bestBeam.generatedTokens());
            if (shared < bestBeam.generatedTokens().size()) {
                List<Long> finalDelta = bestBeam.generatedTokens().subList(shared, bestBeam.generatedTokens().size());
                String deltaText = tokenManager.detokenize(new ArrayList<>(finalDelta));
                if (deltaText != null && !deltaText.isEmpty()) {
                    state.appendText(deltaText);
                }
            }
            committedTokens = new ArrayList<>(bestBeam.generatedTokens());
            bestText = tokenManager.detokenize(bestBeam.generatedTokens());
        }

        return bestText == null ? "" : bestText.trim();
    }

    private ModelInput buildModelInput(List<Long> context, BeamState beam) {
        boolean useCache = model.getConfig().isUseCache();
        boolean isFirstRun = !beam.cacheInitialized();
        int seqLen = (!useCache || isFirstRun) ? context.size() : 1;

        List<Long> inputIds;
        if (seqLen == 1) {
            long token = context.get(context.size() - 1);
            inputIds = Collections.singletonList(token);
        } else {
            inputIds = context;
        }

        int pastSeqLen = isFirstRun ? 0 : Math.max(0, context.size() - 1);
        Map<String, Object> cache = beam.cache() == null ? Collections.emptyMap() : beam.cache();
        return new ModelInput(inputIds, cache, isFirstRun, pastSeqLen);
    }

    private List<Long> buildContext(List<Long> promptTokens, BeamState beam) {
        if (beam.generatedTokens().isEmpty()) {
            return new ArrayList<>(promptTokens);
        }
        List<Long> context = new ArrayList<>(promptTokens.size() + beam.generatedTokens().size());
        context.addAll(promptTokens);
        context.addAll(beam.generatedTokens());
        return context;
    }

    private List<TokenScore> selectCandidates(
        float[] logits,
        List<Long> context,
        SamplingConfig samplingConfig,
        int maxCandidates
    ) {
        if (logits == null || logits.length == 0) {
            return List.of();
        }

        double[] adjusted = SamplingUtils.toDoubleArray(logits);
        SamplingUtils.applyRepetitionPenalty(adjusted, context, samplingConfig.getRepetitionPenalty());
        float temperature = samplingConfig.getTemperature() <= 0f ? 1f : samplingConfig.getTemperature();
        SamplingUtils.applyTemperature(adjusted, temperature);

        int candidateCount = Math.max(maxCandidates, samplingConfig.getTopK());
        List<SamplingUtils.Candidate> rawCandidates = SamplingUtils.buildCandidates(
            adjusted,
            candidateCount,
            samplingConfig.getTopP()
        );

        double logSumExp = logSumExp(adjusted);
        if (rawCandidates.isEmpty()) {
            List<Token> fallbackTokens = samplingStrategy.selectTopNTokens(
                logits,
                context,
                samplingConfig,
                Math.max(1, maxCandidates)
            );
            if (fallbackTokens.isEmpty()) {
                return List.of();
            }
            List<TokenScore> fallbackScores = new ArrayList<>();
            for (Token token : fallbackTokens) {
                int tokenId = (int) token.getId();
                if (tokenId < 0 || tokenId >= adjusted.length) {
                    continue;
                }
                fallbackScores.add(new TokenScore(token.getId(), adjusted[tokenId] - logSumExp));
                if (fallbackScores.size() >= maxCandidates) {
                    break;
                }
            }
            return fallbackScores;
        }

        List<TokenScore> scores = new ArrayList<>();
        for (SamplingUtils.Candidate candidate : rawCandidates) {
            if (scores.size() >= maxCandidates) {
                break;
            }
            int tokenId = candidate.tokenId();
            double logProb = adjusted[tokenId] - logSumExp;
            scores.add(new TokenScore(tokenId, logProb));
        }
        return scores;
    }

    private double logSumExp(double[] values) {
        double max = Double.NEGATIVE_INFINITY;
        for (double value : values) {
            if (value > max) {
                max = value;
            }
        }
        if (Double.isInfinite(max)) {
            return max;
        }
        double sum = 0d;
        for (double value : values) {
            sum += Math.exp(value - max);
        }
        return max + Math.log(Math.max(sum, 1e-12));
    }

    private List<Map<String, Object>> allocateCaches(Map<String, Object> baseCache, int count) {
        if (count <= 0) {
            cacheManager.cleanupCache(baseCache);
            return List.of();
        }
        List<Map<String, Object>> caches = new ArrayList<>(count);
        caches.add(baseCache);
        for (int i = 1; i < count; i++) {
            caches.add(cacheManager.deepCopyCache(baseCache));
        }
        return caches;
    }

    private void logStepSummary(int step, List<BeamState> beams, int stableLen, List<Long> committedTokens) {
        if (beams == null || beams.isEmpty()) {
            logger.debug("🔚 Step #{}: 所有束均完成或无可用束。", step + 1);
            return;
        }
        BeamState best = beams.get(0);
        String bestPreview = previewIds(best.generatedTokens(), 10);
        logger.debug("""

┌─ Beam Step #{}
│ Beams: {}
│ Stable Prefix: {} tokens (committed total: {})
│ Best Beam Score: {:.4f}
│ Best Beam IDs: {}
└──────────────────────────────────────────────
""",
            step + 1,
            beams.size(),
            stableLen,
            committedTokens.size(),
            best.normalizedScore(1.0f),
            bestPreview
        );
    }

    private String previewIds(List<Long> tokens, int limit) {
        if (tokens == null || tokens.isEmpty()) {
            return "[]";
        }
        int size = Math.min(tokens.size(), limit);
        List<Long> preview = tokens.subList(0, size);
        String text = preview.toString();
        if (tokens.size() > limit) {
            text = text.substring(0, text.length() - 1) + ", ...]";
        }
        return text;
    }

    private int commonPrefixLength(List<Long> a, List<Long> b) {
        int max = Math.min(a.size(), b.size());
        for (int i = 0; i < max; i++) {
            if (!Objects.equals(a.get(i), b.get(i))) {
                return i;
            }
        }
        return max;
    }

    private int stablePrefixLength(List<BeamState> beams) {
        if (beams == null || beams.isEmpty()) {
            return 0;
        }
        List<Long> reference = beams.get(0).generatedTokens();
        int prefixLen = reference.size();
        for (int i = 1; i < beams.size(); i++) {
            prefixLen = Math.min(prefixLen, commonPrefixLength(reference, beams.get(i).generatedTokens()));
            if (prefixLen == 0) {
                break;
            }
        }
        return prefixLen;
    }

    private record TokenScore(long tokenId, double logProb) {
    }

    private static final class BeamState {
        private final List<Long> generatedTokens;
        private final double logProb;
        private final Map<String, Object> cache;
        private final boolean cacheInitialized;
        private final boolean finished;

        private BeamState(
            List<Long> generatedTokens,
            double logProb,
            Map<String, Object> cache,
            boolean cacheInitialized,
            boolean finished
        ) {
            this.generatedTokens = generatedTokens;
            this.logProb = logProb;
            this.cache = cache;
            this.cacheInitialized = cacheInitialized;
            this.finished = finished;
        }

        static BeamState initial(Map<String, Object> cache) {
            return new BeamState(new ArrayList<>(), 0.0, cache, false, false);
        }

        BeamState extend(TokenScore tokenScore, Map<String, Object> cacheSnapshot, boolean eos) {
            List<Long> nextTokens = new ArrayList<>(generatedTokens.size() + 1);
            nextTokens.addAll(generatedTokens);
            nextTokens.add(tokenScore.tokenId());
            return new BeamState(
                nextTokens,
                logProb + tokenScore.logProb(),
                cacheSnapshot,
                true,
                eos
            );
        }

        double normalizedScore(float lengthPenalty) {
            int length = Math.max(1, generatedTokens.size());
            double penalty = Math.pow((5.0 + length) / 6.0, Math.max(1e-6, lengthPenalty));
            return logProb / penalty;
        }

        List<Long> generatedTokens() {
            return generatedTokens;
        }

        Map<String, Object> cache() {
            return cache;
        }

        boolean cacheInitialized() {
            return cacheInitialized;
        }

        boolean isFinished() {
            return finished;
        }
    }

    /**
     * 简化的回调状态管理类
     */
    private static class SimpleCallbackState {
        private final GenerationCallback callback;
        private final StringBuilder fullText = new StringBuilder();

        SimpleCallbackState(GenerationCallback callback) {
            this.callback = callback;
        }

        void appendText(String text) {
            if (text != null && !text.isEmpty()) {
                fullText.append(text);
                if (callback != null) {
                    callback.onTokenGenerated(text);
                }
            }
        }

        void finish() {
            if (callback != null) {
                callback.onComplete(fullText.toString());
            }
        }

        void error(Exception e) {
            if (callback != null) {
                callback.onError(e);
            }
        }
    }
}


